#!/usr/bin/env python3

import json
import os
import subprocess

from argparse import ArgumentParser
from dataclasses import dataclass

TOOLS = ('blockdev', 'dd', 'e2image', 'sfdisk')


@dataclass
class PartInfo:
    start: int
    size: int
    thetype: str



def partition_target(dst, esp, root):
    ptbl = f"""
    label: gpt
    1: start={esp.start}, size={esp.size}, type=uefi
    2: start={root.start}, size={root.size}, type={root.thetype}
    """
    cmd = 'sfdisk --no-reread --wipe always --wipe-partition always'.split()
    cmd.append(dst)
    subprocess.run(cmd, input=ptbl, check=True, encoding='utf-8')


def copy_rootfs(src, dst, root_info):
    root_offset = root_info.start * 512
    cmd = f'e2image -rap -c -o {root_offset} -O {root_offset}'.split()
    cmd += [src, dst]
    subprocess.check_call(cmd)


def copy_esp(src, dst, esp_info):
    cmd = [
        'dd',
        f'if={src}',
        f'of={dst}',
        'bs=512',
        f'skip={esp_info.start}',
        f'seek={esp_info.start}',
        f'count={esp_info.size}',
    ]
    subprocess.check_call(cmd, encoding='utf-8')


def reread_parition_table(dst):
    cmd = 'sudo blockdev --rereadpt'.split()
    cmd += [dst]
    subprocess.check_call(cmd, encoding='utf-8')


def verify_image(src):
    cmd = 'sfdisk -l --json'.split()
    cmd.append(src)
    out = subprocess.check_output(cmd, encoding='utf-8')
    img_info = json.loads(out)['partitiontable']
    partitions = img_info['partitions']
    assert len(partitions) == 2

    def to_partinfo(p):
        return PartInfo(start=int(p['start']),
                        size=int(p['size']),
                        thetype=p['type'])

    root = to_partinfo(partitions[1])
    esp = to_partinfo(partitions[0])
    return esp, root


def disk_size_sect(dst):
    cmd = ['blockdev', '--getsz', dst]
    out = subprocess.check_output(cmd, encoding='utf-8')
    return int(out.strip())


def is_target_big_enough(esp, root, dest_size):
    assert dest_size >= esp.start + esp.size
    assert dest_size > root.start
    assert dest_size >= root.start + root.size


def adjust_root_size(esp, root, dest_size):
    esp_offset = esp.start
    old_root_size = root.size
    # XXX: there must be a 2nd copy of GPT at the end of the drive
    new_root_size = dest_size - esp.size - 2*esp_offset #
    assert new_root_size >= old_root_size
    if new_root_size > root.size:
        root.size = new_root_size
    return root.size > old_root_size


def is_in_PATH(tool):
    cmd = [tool, '--version']
    try:
        subprocess.check_call(cmd)
        return True
    except FileNotFoundError as err:
        return False
    except:
        # runs, but don't understand '--version'
        return True


def adjust_PATH():
    path = os.environ['PATH'].split(':')
    if not is_in_PATH('blockdev'):
        path += ['/sbin', '/usr/sbin']
        os.environ['PATH'] = ':'.join(path)


def check_tools():
    for prog in TOOLS:
        assert is_in_PATH(prog)


def deploy(src, dst):
    adjust_PATH()
    check_tools()
    esp, root = verify_image(src)
    dest_size = disk_size_sect(dst)
    adjust_root_size(esp, root, dest_size)
    partition_target(dst, esp, root)
    copy_esp(src, dst, esp)
    copy_rootfs(src, dst, root)
    reread_parition_table(dst)


def main():
    parser = ArgumentParser()
    parser.add_argument('src', help='Image to deploy')
    parser.add_argument('dst',
        help='Destination device, ALL DATA WILL BE DESTROYED')
    args = parser.parse_args()
    deploy(args.src, args.dst)


if __name__ == '__main__':
    main()
